import os
import numpy as np
import logging
import shutil
from tqdm import tqdm
from collections import Counter, OrderedDict
import dominate
from dominate.tags import meta, h3, table, tr, td, p, a, img, br
import time
import cv2

def get_subdirs(data_dir):
    for i, (root, dirs, files) in enumerate(os.walk(data_dir)): 
        if i == 0:
            sub_dirs = dirs

    return sub_dirs

def set_logger(log_path, logging_name):
    """Sets the logger to log info in terminal and file `log_path`.

    In general, it is useful to have a logger so that every output to the terminal is saved
    in a permanent file. Here we save it to `model_dir/train.log`.

    Example:
    ```
    logging.info("Starting training...")
    ```

    Args:
        log_path: (string) where to log
    """
    _logger = logging.getLogger(logging_name)
    _logger.setLevel(logging.INFO)

    fmt = logging.Formatter('[%(asctime)s]:%(name)s: %(message)s', '%H:%M:%S')

    class TqdmHandler(logging.StreamHandler):
        def __init__(self, formatter):
            logging.StreamHandler.__init__(self)
            self.setFormatter(formatter)

        def emit(self, record):
            msg = self.format(record)
            tqdm.write(msg)

    file_handler = logging.FileHandler(log_path)
    file_handler.setFormatter(fmt)
    _logger.addHandler(file_handler)
    _logger.addHandler(TqdmHandler(fmt))

def getDirectoryList(path=None, extension='.txt'):
    """
    Find subdirectories that contains specific files
    """
    directoryList = []
    # print('Finding all subdirectories that contains %s files.' %extension)
    #return nothing if path is a file
    if os.path.isfile(path):
        return []

    #add dir to directorylist if it contains .txt files
    if len([f for f in os.listdir(path) if f.endswith(extension)])>0:
        directoryList.append(path)

    for d in os.listdir(path):
        new_path = os.path.join(path, d)
        if os.path.isdir(new_path):
            directoryList += getDirectoryList(new_path, extension)

    return directoryList

def del_useless_folders(path=None, remain_exts='.pth'):
    """
    Deleting all subdirectories that not contain specific files.
    """
    assert isinstance(remain_exts, (str, list, tuple))
    print('Deleting all subdirectories that not contain %s files.' %remain_exts)
    subfolders = [ f.path for f in os.scandir(path) if f.is_dir() ]
    if isinstance(remain_exts, str):
        remain_list = getDirectoryList(path, remain_exts)
    else:
        remain_list = []
        for extension in remain_exts:
            remain_list += getDirectoryList(path, extension)
        remain_list = list(set(remain_list))
    
    del_list = [i for i in subfolders if i not in remain_list]
    if len(del_list) > 0:
        for del_folder in del_list:
            shutil.rmtree(del_folder)
            print('Folder %s has been deleted.' %del_folder)

def remove_dataparallel(state_dict):
    """Remover data parallel when loading state dict
    Args:
        state_dict: saved state dict with data parallel

    Returns:
        new state dict without `module`
    Raises:
        IOError: An error occurred accessing state_dice object.
    """
    new_state_dict = OrderedDict()
    for k,v in state_dict.items():
        if k[:7] == "module.": #remove module.
            name = k[7:]
        else:
            name = k
        new_state_dict[name] = v

    return new_state_dict

def denormalization(x, mean=[0.485, 0.456, 0.406], std=[0.229, 0.224, 0.225]):
    # mean = np.array([0.485, 0.456, 0.406])
    # std = np.array([0.229, 0.224, 0.225])
    mean = [float(x) for x in mean]
    std = [float(x) for x in std]
    x = (((x.transpose(1, 2, 0) * std) + mean) * 255.).astype(np.uint8)
    # x = (x.transpose(1, 2, 0) * 255.).astype(np.uint8)
    return x


def segImage(I, S, *args):
    # SEGIMAGE Given an image and its segmentation, plot the segment boundaries
    # Input
    # I               input image
    # S              input segmentation
    # varargin   width: width of the boundary
    #                 color: color of the boundary
    # Output
    # I               output image with segment boundaries
    debug = 0
    varargin = args
    boundaryWidth = 1
    boundaryColor = 'red'
    if debug:
        t0 = time.time()
    # Parse the optional inputs.
    if (len(varargin) % 2 != 0):
        print('Extra Parameters passed to the function must be passed in pairs.')

    parameterCount = int(len(varargin) / 2)

    for parameterIndex in range(parameterCount):
        parameterName = varargin[parameterIndex * 2 - 1]
        parameterValue = varargin[parameterIndex * 2]
        switcher = str.lower(parameterName)
        if switcher == 'boundarycolor':
            boundaryColor = parameterValue
        elif switcher == 'boundarywidth':
            boundaryWidth = parameterValue
        else:
            print('Invalid parameter')

    # S = np.array(S, dtype = np.int16)
    S = np.int16(S)
    # S = im2db(S)
    [cy, cx] = np.gradient(S)
    if debug:
        t1 = time.time()
        print('gradient take time = {}'.format(t1 - t0))
    ccc = np.where((abs(cx) + abs(cy)) > 0, 1, 0)
    ccc1 = np.where(ccc > 0, 0, 1)

    # ccc = ccc.astype(int)
    # ccc1 = ccc1.astype(int)
    ccc = np.uint8(ccc)
    ccc1 = np.uint8(ccc1)
    if boundaryWidth > 1:
        boundaryWidth = np.ceil(boundaryWidth)
        dilateWindow = np.ones(boundaryWidth, boundaryWidth)
        ccc = cv2.dilate(ccc, dilateWindow, iterations=1)  # imdilate(ccc,dilateWindow)
    elif boundaryWidth < 1:
        print('boundaryWidth has been reset to 1.')

    if boundaryColor == 'red':
        I[:, :, 0] = np.maximum(I[:, :, 0], ccc)
        I[:, :, 1] = np.minimum(I[:, :, 1], ccc1)
        I[:, :, 2] = np.minimum(I[:, :, 2], ccc1)
    elif boundaryColor == 'black':
        I[:, :, 0] = np.minimum(I[:, :, 0], ccc1)
        I[:, :, 1] = np.minimum(I[:, :, 1], ccc1)
        I[:, :, 2] = np.minimum(I[:, :, 2], ccc1)
    else:
        print('Does not recognize boundaryColor other than red and black')

        I[:, :, 0] = np.maximum(I[:, :, 0], ccc1)
        I[:, :, 1] = np.maximum(I[:, :, 1], ccc1)
        I[:, :, 2] = np.maximum(I[:, :, 2], ccc1)
    if debug:
        t2 = time.time()
        print('segImg take time = {}'.format(t2 - t1))

    # pyplot.imshow(I)
    # pyplot.show()

    #    amin = np.min(I)
    #    amax = np.max(I)
    #    I = (I - amin)/amax
    return I

class Viz_HTML:
    """This HTML class allows us to save images and write texts into a single HTML file.
     It consists of functions such as <add_header> (add a text header to the HTML file),
     <add_images> (add a row of images to the HTML file), and <save> (save the HTML to the disk).
     It is based on Python library 'dominate', a Python library for creating and manipulating HTML documents using a DOM API.
    """

    def __init__(self, web_dir, title, refresh=0):
        """Initialize the HTML classes
        Parameters:
            web_dir (str) -- a directory that stores the webpage. HTML file will be created at <web_dir>/index.html; images will be saved at <web_dir/images/
            title (str)   -- the webpage name
            refresh (int) -- how often the website refresh itself; if 0; no refreshing
        """
        self.title = title
        self.web_dir = web_dir
        self.img_dir = os.path.join(self.web_dir, 'images')
        if not os.path.exists(self.web_dir):
            os.makedirs(self.web_dir)
        if not os.path.exists(self.img_dir):
            os.makedirs(self.img_dir)

        self.doc = dominate.document(title=title)
        if refresh > 0:
            with self.doc.head:
                meta(http_equiv="refresh", content=str(refresh))

    def get_image_dir(self):
        """Return the directory that stores images"""
        return self.img_dir

    def add_header(self, text):
        """Insert a header to the HTML file
        Parameters:
            text (str) -- the header text
        """
        with self.doc:
            h3(text)

    def add_images(self, ims, txts, links, width=400):
        """add images to the HTML file
        Parameters:
            ims (str list)   -- a list of image paths
            txts (str list)  -- a list of image names shown on the website
            links (str list) --  a list of hyperref links; when you click an image, it will redirect you to a new page
        """
        self.t = table(border=1, style="table-layout: fixed;")  # Insert a table
        self.doc.add(self.t)
        with self.t:
            with tr():
                for im, txt, link in zip(ims, txts, links):
                    with td(style="word-wrap: break-word;", halign="center", valign="top"):
                        with p():
                            with a(href=os.path.join('images', link)):
                                img(style="width:%dpx" % width, src=os.path.join('images', im))
                            br()
                            p(txt)

    def save(self):
        """save the current content to the HMTL file"""
        html_file = '%s/seg_viz.html' % self.web_dir
        f = open(html_file, 'wt')
        f.write(self.doc.render())
        f.close()